// SeedableRng is not an unused import, it's import-ant as a trait, dammit!
use rand::Rng;
#[cfg(test)]
use rand::SeedableRng;

use num_bigint::BigUint;
#[cfg(test)]
use num_rational::Ratio;
use num_traits::identities::{One, Zero};

use std::iter::repeat;

use crate::types::BinaryFraction;

#[cfg(test)]
lazy_static! {
    // use One::one, Zero::zero where possible
    static ref ONE_R: Ratio<BigUint> = Ratio::from_integer(BigUint::parse_bytes(b"1", 2).unwrap());
}

pub fn xor_blocks(block1: &[u8], block2: &[u8]) -> Vec<u8> {
    assert!(block1.len() == block2.len());
    let mut out: Vec<u8> = Vec::with_capacity(block1.len());
    block1.iter().zip(block2.iter()).for_each(|(x1, x2)| {
        out.push(*x1 ^ *x2);
    });
    out
}

/********************
 * BigNum functions *
 ********************/
/// input:  bytestring (non empty)
/// output: a number encoding the bytestring suffixed with a terminating "1" if it ends in a "0" (which cannot be distinguished from end of message otherwise)
/// @note used only to test BinaryFraction against Ratio<BigUInt>
#[cfg(test)]
pub fn bytes_to_bigrational_between_0_and_1(input: &[u8]) -> Ratio<BigUint> {
    let endzero = input[input.len() - 1] & 1 == 0;

    let mut extended_input = input.to_vec();

    if endzero {
        extended_input.extend(vec![0x80 as u8]);
    }

    let extended_input = extended_input;

    let denom_numzerobytes = extended_input.len();
    let denom_leadingbyte: u8 = 0x01;

    let denom_bytes = vec![denom_leadingbyte as u8]
        .into_iter()
        .chain(repeat(0x00 as u8).take(denom_numzerobytes))
        .collect::<Vec<u8>>();

    let numer: BigUint = BigUint::from_bytes_be(&extended_input);
    let denom: BigUint = BigUint::from_bytes_be(&denom_bytes);

    Ratio::new(numer, denom)
}

pub fn bytes_to_binaryfraction_between_0_and_1(input: &[u8]) -> BinaryFraction {
    let endzero = input[input.len() - 1] & 1 == 0;
    let mut extended_input = input.to_vec();
    if endzero {
        extended_input.extend(vec![0x80_u8]);
    }
    let extended_input = extended_input;

    BinaryFraction::new(
        BigUint::from_bytes_be(&extended_input),
        extended_input.len() as i64 * 8,
    )
}

pub fn shl(r: BinaryFraction, z: u32) -> BinaryFraction {
    BinaryFraction::new(r.k << z, r.f)
}

pub fn shl8(r: BinaryFraction) -> BinaryFraction {
    BinaryFraction::new(r.k << 8, r.f)
}

pub fn shr(r: BinaryFraction, z: u32) -> BinaryFraction {
    BinaryFraction::new(r.k, r.f + (z as i64))
}

/// Find out how many zeroes after the binary point to the first 1 (e.g. 0.00001 -> 4)
/// @input r: a number between [0 and 1), that has a finite binary expansion
#[cfg(test)]
pub fn leading_zeroes(r: &Ratio<BigUint>) -> Option<u32> {
    // could even forgo this check at runtime
    if r.is_zero() || *r >= One::one() {
        None
    } else {
        Some((r.denom().bits() - r.numer().bits() - 1) as u32)
    }
}

pub fn binaryfraction_leading_zeroes(r: &BinaryFraction) -> Option<u32> {
    // could even forgo this check at runtime
    if r.is_zero() || *r >= One::one() {
        None
    } else {
        Some((r.f - r.k.bits() as i64) as u32)
    }
}

#[derive(PartialEq, Eq, Debug)]
pub enum BinaryFractionLongestCommonPrefix {
    Same,
    Invalid,
    Length { bits: u32 },
}

pub fn binaryfractions_length_common_prefix(
    r: &BinaryFraction,
    s: &BinaryFraction,
) -> BinaryFractionLongestCommonPrefix {
    if (*r >= One::one()) && (*s >= One::one()) {
        BinaryFractionLongestCommonPrefix::Invalid
    } else if r == s {
        BinaryFractionLongestCommonPrefix::Same
    } else if r.f == s.f {
        match compare_same_f(r.f, &r.k, &s.k) {
            None => BinaryFractionLongestCommonPrefix::Same,
            Some(bits) => BinaryFractionLongestCommonPrefix::Length { bits: bits as u32 },
        }
    } else if r.f > s.f {
        let s_shl = std::ops::Shl::shl(&s.k, r.f - s.f);
        match compare_same_f(r.f, &r.k, &s_shl) {
            None => BinaryFractionLongestCommonPrefix::Same,
            Some(bits) => BinaryFractionLongestCommonPrefix::Length { bits: bits as u32 },
        }
    } else {
        let r_shl = std::ops::Shl::shl(&r.k, s.f - r.f);
        match compare_same_f(s.f, &r_shl, &s.k) {
            None => BinaryFractionLongestCommonPrefix::Same,
            Some(bits) => BinaryFractionLongestCommonPrefix::Length { bits: bits as u32 },
        }
    }
}

// TODO deduplicate code
fn lzbits(f: u32, u64s_r: &[u64]) -> Option<u32> {
    let n = u64s_r.len();

    if f < n as u32 * 64 {
        let mut result = 0;

        for j in 1..=n {
            let xr = u64s_r.get(n - j).unwrap();
            let differing_digits = xr;
            if *differing_digits == 0 {
                result += if j == 1 { f % 64 } else { 64 };
            } else {
                for i in 0u8..64 {
                    let bit = (differing_digits >> (63 - i)) & 1u64;
                    if bit == 0 {
                        result += 1;
                    } else {
                        if j == 1 {
                            let x = (f % 64) as i32;
                            if (result as i32 - (64 - x)) < 0 {
                                return Some(result + x as u32);
                            } else {
                                return Some((result as i32 - (64 - x)) as u32);
                            }
                        }
                        return Some(result);
                    }
                }
            }
        }
    }

    let excess_bits = (f % 64) as u32;
    let excess_bytes = ((f - (n as u32 * 64)) / 64) as u32;
    let extra = excess_bytes * 64 + excess_bits;

    let mut result = 0;

    for j in 1..=n {
        let xr = u64s_r.get(n - j).unwrap();
        let differing_digits = xr;
        if *differing_digits == 0 {
            result += 64;
        } else {
            for i in 0u8..64 {
                let bit = (differing_digits >> (63 - i)) & 1u64;
                if bit == 0 {
                    result += 1;
                } else {
                    return Some(result + extra);
                }
            }
        }
    }

    None
}

// Candidate for optimization? And should be better tested / verified
fn compare_same_f(f: i64, r: &BigUint, s: &BigUint) -> Option<u32> {
    let u64s_r = r.to_u64_digits();
    let u64s_s = s.to_u64_digits();

    let f = f as u32;

    if u64s_r.len() == u64s_s.len() {
        let n = u64s_r.len();

        if f < n as u32 * 64 {
            let mut result = 0;

            for j in 1..=n {
                let xr = u64s_r.get(n - j).unwrap();
                let xs = u64s_s.get(n - j).unwrap();
                let differing_digits = xr ^ xs;
                if differing_digits == 0 {
                    result += if j == 1 { f % 64 } else { 64 };
                } else {
                    for i in 0u8..64 {
                        let bit = (differing_digits >> (63 - i)) & 1u64;
                        if bit == 0 {
                            result += 1;
                        } else {
                            if j == 1 {
                                let x = (f % 64) as i32;
                                if (result as i32 - (64 - x)) < 0 {
                                    return Some(result + x as u32);
                                } else {
                                    return Some((result as i32 - (64 - x)) as u32);
                                }
                            }
                            return Some(result);
                        }
                    }
                }
            }
        }
        let excess_bits = (f % 64) as u32;
        let excess_bytes = ((f - (n as u32 * 64)) / 64) as u32;
        let extra = excess_bytes * 64 + excess_bits;

        let mut result = 0;

        for j in 1..=n {
            let xr = u64s_r.get(n - j).unwrap();
            let xs = u64s_s.get(n - j).unwrap();
            let differing_digits = xr ^ xs;
            if differing_digits == 0 {
                result += 64;
            } else {
                for i in 0u8..64 {
                    let bit = (differing_digits >> (63 - i)) & 1u64;
                    if bit == 0 {
                        result += 1;
                    } else {
                        return Some(result + extra);
                    }
                }
            }
        }
        None // should not be reached
    } else if u64s_r.len() > u64s_s.len() {
        lzbits(f, &u64s_r)
    } else {
        lzbits(f, &u64s_s)
    }
}

pub fn binaryfractions_same_prefix(bits: u32, r: &BinaryFraction, s: &BinaryFraction) -> bool {
    shl(r.clone(), bits).trunc_as_biguint() == shl(s.clone(), bits).trunc_as_biguint()
}

/// Recover bytestring from binary fraction
/// @precondition: input is in [0 .. 1)
/// @precondition: input's denominator is a power of 2 (not checked)
#[cfg(test)]
pub fn bigrational_between_0_and_1_to_bytes(input: &Ratio<BigUint>) -> Vec<u8> {
    match input.denom().trailing_zeros() {
        Some(log) => {
            let align_shift: u32 = 8 - log as u32 % 8;
            match leading_zeroes(input) {
                Some(lzbits) => {
                    let lzbytes: u32 = lzbits / 8;
                    let mut appendix = (input.numer() << align_shift).to_bytes_be();
                    let mut result = repeat(0x00 as u8)
                        .take(lzbytes as usize)
                        .collect::<Vec<u8>>();
                    result.append(&mut appendix);
                    result
                }
                None => {
                    panic!(
                        "bigrational_between_0_and_1_to_bytes: input {} out of range",
                        input
                    );
                }
            }
        }
        None => {
            panic!(
                "bigrational_between_0_and_1_to_bytes: input {} out of range",
                input
            );
        }
    }
}

pub fn binaryfraction_between_0_and_1_to_bytes(input: &BinaryFraction) -> Vec<u8> {
    if input.f > 0 {
        let log = input.f;
        let align_shift: u32 = 8 - log as u32 % 8;
        match binaryfraction_leading_zeroes(input) {
            Some(lzbits) => {
                let lzbytes: u32 = lzbits / 8;
                let mut appendix = (input.k.clone() << align_shift).to_bytes_be();
                let mut result = repeat(0x00_u8).take(lzbytes as usize).collect::<Vec<u8>>();
                result.append(&mut appendix);
                result
            }
            None => {
                panic!(
                    "bigrational_between_0_and_1_to_bytes: input {} out of range",
                    input
                );
            }
        }
    } else {
        panic!(
            "binaryfraction_between_0_and_1_to_bytes: input {} out of range",
            input
        );
    }
}

/*****************************
 * num_utils (non-essential) *
 *****************************/

#[cfg(feature = "linalg")]
pub mod num_utils {
    use num_traits::{One, Zero};

    pub fn exp_associative<M: std::clone::Clone, T>(a: M, p: T) -> M
    where
        T: num_traits::One
            + std::hash::Hash
            + num_traits::Zero
            + std::cmp::PartialOrd
            + num_integer::Integer
            + std::ops::Shr<Output = T>
            + Copy
            + std::fmt::Debug,
        M: std::ops::Mul<Output = M>,
    {
        assert!(p >= Zero::zero());
        if p == Zero::zero() {
            unimplemented!();
        // should be M::one()
        } else if p == One::one() {
            a
        } else if p.is_odd() {
            let sub = exp_associative(a.clone(), p - One::one());
            sub * a
        } else {
            let sub = exp_associative(a, p >> One::one());
            sub.clone() * sub
        }
    }

    pub fn lcm<T>(ns: &[T]) -> T
    where
        T: num_traits::One
            + std::hash::Hash
            + std::ops::SubAssign
            + std::cmp::PartialEq
            + std::cmp::Ord
            + std::ops::Div<Output = T>
            + Copy,
    {
        if ns.is_empty() {
            One::one()
        } else if ns.len() == 1 {
            ns[0]
        } else {
            let lcm1 = lcm(&ns[0..ns.len() / 2]);
            let lcm2 = lcm(&ns[ns.len() / 2..]);

            let mut n1 = lcm1;
            let mut n2 = lcm2;

            let gcd = if std::cmp::max(n1, n2).is_one() {
                One::one()
            } else {
                loop {
                    if n1 == n2 {
                        break n1;
                    }
                    if n1 > n2 {
                        n1 -= n2;
                        continue;
                    } else {
                        n2 -= n1;
                        continue;
                    }
                }
            };

            (lcm1 * lcm2) / gcd
        }
    }
}

/**********
 * RANDOM *
 **********/

pub struct Randoms {
    pub rng: rand::rngs::ThreadRng,
}

impl Iterator for Randoms {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        Some(self.rng.gen())
    }
}

#[cfg(test)]
pub struct DeterministicPseudoRandoms {
    rng: rand_chacha::ChaCha8Rng,
}

// See https://rust-random.github.io/book/guide-seeding.html
#[cfg(test)]
impl DeterministicPseudoRandoms {
    // this initialization is not intended to be secure, it's just for testing purposes
    pub fn new(miniseed: u64) -> Self {
        Self {
            rng: rand_chacha::ChaCha8Rng::seed_from_u64(miniseed),
        }
    }
}

#[cfg(test)]
impl Iterator for DeterministicPseudoRandoms {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        return Some(self.rng.gen());
    }
}
